extern "C" {
#include <immintrin.h>
}

#define A(i, j) A[(i) + (j) * (LDA)]
#define B(i, j) B[(i) + (j) * (LDB)]
#define C(i, j) C[(i) + (j) * (LDC)]

static void mm_v5_corner(int M, int N, int K, double *A, int LDA, double *B,
                         int LDB, double *C, int LDC) {
    double tmp = 0;
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            tmp = 0;
            for (int k = 0; k < K; k++) {
                tmp += A(i, k) * B(k, j);
            }
            C(i, j) = tmp;
        }
    }
}
// 重复利用寄存器 4x4分块
// A:[M, K] B:[K, N] C:[M, N]
inline void mm_v5(int M, int N, int K, double *A, int LDA, double *B, int LDB,
                  double *C, int LDC) {
    int m4 = M & -4;
    int n4 = N & -4;
    for (int i = 0; i < m4; i += 4) {
        for (int j = 0; j < n4; j += 4) {
            __m256d c0 = _mm256_setzero_pd();
            __m256d c1 = _mm256_setzero_pd();
            __m256d c2 = _mm256_setzero_pd();
            __m256d c3 = _mm256_setzero_pd();
            for (int k = 0; k < K; k++) {
                __m256d a = _mm256_loadu_pd(
                    &A(i, k)); // a:[A(i,k) A(i+1,k) A(i+2,k) A(i+3,k)]
                __m256d b0 = _mm256_broadcast_sd(
                    &B(k, j)); // b0:[B(k,j) B(k,j) B(k,j) B(k,j)]
                __m256d b1 = _mm256_broadcast_sd(
                    &B(k, j + 1)); // b1:[B(k,j+1) B(k,j+1) B(k,j+1) B(k,j+1)]
                __m256d b2 = _mm256_broadcast_sd(
                    &B(k, j + 2)); // b2:[B(k,j+2) B(k,j+2) B(k,j+2) B(k,j+2)]
                __m256d b3 = _mm256_broadcast_sd(
                    &B(k, j + 3)); // b3:[B(k,j+3) B(k,j+3) B(k,j+3) B(k,j+3)]

                c0 = _mm256_fmadd_pd(
                    a, b0,
                    c0); // c0:[C(i,j) C(i+1,j) C(i+2,j) C(i+3,j)]
                c1 = _mm256_fmadd_pd(
                    a, b1,
                    c1); // c1:[C(i,j+1) C(i+1,j+1) C(i+2,j+1) C(i+3,j+1)]
                c2 = _mm256_fmadd_pd(
                    a, b2,
                    c2); // c2:[C(i,j+2) C(i+1,j+2) C(i+2,j+2) C(i+3,j+2)]
                c3 = _mm256_fmadd_pd(
                    a, b3,
                    c3); // c3:[C(i,j+3) C(i+1,j+3) C(i+2,j+3) C(i+3,j+3)]
            }
            _mm256_storeu_pd(&C(i, j), c0);
            _mm256_storeu_pd(&C(i, j + 1), c1);
            _mm256_storeu_pd(&C(i, j + 2), c2);
            _mm256_storeu_pd(&C(i, j + 3), c3);
        }
    }
    if (M == m4 && N == n4) {
        return;
    }
    if (M != m4) {
        mm_v5_corner(M - m4, N, K, &A(m4, 0), LDA, &B(0, 0), LDB, &C(m4, 0),
                     LDC);
    }
    if (N != n4) {
        mm_v5_corner(m4, N - n4, K, &A(0, 0), LDA, &B(0, n4), LDB, &C(0, n4),
                     LDC);
    }
}
#undef A
#undef B
#undef C
